#!/usr/bin/env python
# -*- coding:utf-8 -*-

import select
import socket
import sys, os
import time

from errno import EALREADY, EINPROGRESS, EWOULDBLOCK, ECONNRESET, \
     ECONNREFUSED, ENOTCONN, ESHUTDOWN, EINTR, EISCONN, errorcode

from collections import deque

import log

#######################################################################

class   JujuTransporter( object ) :
    RECORD_END_MARK      = '\n'
    RECORD_END_MARK_SIZE = 1

    IN_BUFFER_SIZE = 65536

    def __init__(self, address, jucore=None, sock=None) :
        self.address = address
        self.jucore = jucore
        self.socket = sock

        if self.socket :
            self.fileno = self.socket.fileno()
            self.connected = True
        else :
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.setblocking(0)
            #if fcntl and hasattr(fcntl, 'FD_CLOEXEC'):
            #    old = fcntl.fcntl(self.socket.fileno(), fcntl.F_GETFD)
            #    fcntl.fcntl(self.socket.fileno(), fcntl.F_SETFD, old | fcntl.FD_CLOEXEC)
            self.fileno = self.socket.fileno()
            self.connected = False

        if jucore :
            jucore.addChannel(self)

        self.in_buffer = ''
        self.collecting = []

        self.out_buffer = ''
        self.producer_fifo = Fifo()
        # __init__()


    def readable(self) :
        if self.connected :
            return True
        else :
            return False
        # readable()

    def writable(self) :
        if self.connected :
            if self.out_buffer :
                return True
            if not self.producer_fifo.empty() :
                return True
        else :
            if self.out_buffer :
                self.reconnect()
                if self.connected : return True
            if not self.producer_fifo.empty() :
                self.reconnect()
                if self.connected : return True

        return False
        # writable()

    ###################################################################

    def OnRecvRecord(self, record):
        raise NotImplementedError, "must be implemented in subclass"

    def OnPreSendRecord(self) :
        pass
    def OnPostSendRecord(self) :
        pass

    def OnConnectionLost(self) :
        pass

    def push(self, data) :
        if not data.endswith(self.RECORD_END_MARK) :
            data = data + self.RECORD_END_MARK
        self.producer_fifo.push(SimpleProducer(data))
        # push()
    def pushWithProducer(self, producer) :
        self.producer_fifo.push(producer)
        # pushWithProducer()
    def countProducers(self) :
        return  len(self.producer_fifo)
        # countProducers()

    ###################################################################

    def collectIncoming(self, data) :
        self.collecting.append(data)
        # collectIncoming()

    def OnTerminatorFound(self) :
        data = ''.join(self.collecting)
        self.OnRecvRecord(data)
        self.collecting = []
        # OnTerminatorFound()

    def handleReadEvent(self) :
        try:
            data = self.socket.recv (self.IN_BUFFER_SIZE)
        except socket.error, why:
            log.error('socket READ error: %s', why)
            self.handleError()
            return

        if len(data) < 1 :
            log.error('socket READ nothing, pipe broken.')
            self.handleCloseEvent()
            return
        
        self.in_buffer = self.in_buffer + data

        # Continue to search for terminator in in_buffer,
        # while calling self.collectIncoming.  The while loop
        # is necessary because we might read several data+terminator
        # combos with a single recv.

        terminator = self.RECORD_END_MARK
        terminator_len = len(terminator)
        while self.in_buffer:
            # 3 cases:
            # 1) end of buffer matches terminator exactly:
            #    collect data, transition
            # 2) end of buffer matches some prefix:
            #    collect data to the prefix
            # 3) end of buffer does not match any prefix:
            #    collect data
            index = self.in_buffer.find(terminator)
            if index != -1:
                # we found the terminator
                if index > 0:
                    # don't bother reporting the empty string (source of subtle bugs)
                    self.collectIncoming (self.in_buffer[:index])
                self.in_buffer = self.in_buffer[index+terminator_len:]
                # This does the Right Thing if the terminator is changed here.
                self.OnTerminatorFound()
            else:
                # check for a prefix of the terminator
                index = find_prefix_at_end (self.in_buffer, terminator)
                if index:
                    if index != len(self.in_buffer) :
                        # we found a prefix, collect up to the prefix
                        self.collectIncoming (self.in_buffer[:-index])
                        self.in_buffer = self.in_buffer[-index:]
                    break
                else:
                    # no prefix, collect it all
                    self.collectIncoming (self.in_buffer)
                    self.in_buffer = ''

        # handleReadEvent()

    ###################################################################

    def handleWriteEvent(self) :
        self.sendData()
        # handleWriteEvent()

    def sendData(self) :
        if not self.out_buffer :
            self.refillOutBuffer()
            if not self.out_buffer :
                return False

        if self.out_buffer and self.connected :
            try :
                num_sent = self.socket.send(self.out_buffer)
                if num_sent :
                    self.out_buffer = self.out_buffer[num_sent:]
            except socket.error, why :
                self.handleError()
                return
        # sendData()

    def refillOutBuffer(self) :
        while True :
            if self.producer_fifo.empty() :
                return False
            p = self.producer_fifo.first()
            data = p.more()
            if data :
                self.out_buffer = self.out_buffer + data
                self.OnPreSendRecord()
                return True
            else :
                self.producer_fifo.pop()
                self.OnPostSendRecord()
                continue
        return True
        # refillOutBuffer()

    ###################################################################

    def handleExceptionEvent(self) :
        raise NotImplementedError
        # handleExceptionEvent()

    ###################################################################

    def handleError(self):
        nil, t, v, tbinfo = compact_traceback()

        # sometimes a user repr method will crash.
        try:
            self_repr = repr(self)
        except:
            self_repr = '<__repr__(self) failed for object at %0x>' % id(self)

        log.error(
            'uncaptured python exception, closing channel %s (%s:%s %s)' % (
                self_repr,
                t,
                v,
                tbinfo
                )
            )
        self.handleCloseEvent()
        # handleError()

    def handleCloseEvent(self) :
        if self.jucore :
            # this should work beofre socket.close
            # cause we need socket.fileno()
            self.jucore.delChannel(self)

        self.socket.close()
        self.connected = False

        self.OnConnectionLost()
        # handleCloseEvent()

    ###################################################################

    def reconnect(self) :
        if self.connected : return

        errno = self.socket.connect_ex(self.address)
        if errno in (0, EISCONN) :
            self.connected = True
        elif errno in (EINPROGRESS, EALREADY, EWOULDBLOCK):
            self.connected = False
        else :
            self.connected = False
            log.error('fail to connect to "%s" -- [%d] %s', self.address, errno, errorcode[errno])
            #raise socket.error, (errno, errorcode[errno])
        # reconnect()

    # class JujuTransporter

class   JujuListener( object ) :
    def __init__(self, port, jucore=None) :
        self.port = port
        self.address = ('', port)
        self.jucore = jucore

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.setblocking(0)
        #if fcntl and hasattr(fcntl, 'FD_CLOEXEC'):
        #    old = fcntl.fcntl(self.socket.fileno(), fcntl.F_GETFD)
        #    fcntl.fcntl(self.socket.fileno(), fcntl.F_SETFD, old | fcntl.FD_CLOEXEC)
        if 'posix' == os.name :
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.socket.bind( self.address )
        self.socket.listen( 5 )

        if self.socket :
            self.fileno = self.socket.fileno()
            self.connected = True
        else :
            self.fileno = 0
            self.connected = False

        if jucore :
            jucore.addChannel(self)
        # __init__()

    def handleReadEvent(self) :
        connect, address = self.socket.accept()
        JujuTransporter(address, jucore, connect)
        # handleReadEvent()

    def readable(self) :
        if self.connected :
            return True
        else :
            return False
        # readable()

    def writable(self) :
        return False
        # writable()
    # class JujuListener

class   SimpleProducer( object ) :
    def __init__(self, data, buffer_size=65536) :
        self.data = data
        self.buffer_size = buffer_size
        # __init__()

    def more(self) :
        result = self.data
        self.data = ''

        return result
        # more()

    # class SimpleProducer

class   Fifo:
    def __init__ (self, list=None):
        if not list:
            self.list = deque()
        else:
            self.list = deque(list)

    def __len__ (self):
        return len(self.list)

    def empty (self):
        return not self.list

    def first (self):
        return self.list[0]

    def push (self, data):
        self.list.append(data)

    def pop (self):
        if self.list:
            return (1, self.list.popleft())
        else:
            return (0, None)


def     compact_traceback():
    t, v, tb = sys.exc_info()
    tbinfo = []
    assert tb # Must have a traceback
    while tb:
        tbinfo.append((
            tb.tb_frame.f_code.co_filename,
            tb.tb_frame.f_code.co_name,
            str(tb.tb_lineno)
            ))
        tb = tb.tb_next

    # just to be safe
    del tb

    file, function, line = tbinfo[-1]
    info = ' '.join(['[%s|%s|%s]' % x for x in tbinfo])
    return (file, function, line), t, v, info

# Given 'haystack', see if any prefix of 'needle' is at its end.  This
# assumes an exact match has already been checked.  Return the number of
# characters matched.
# for example:
# f_p_a_e ("qwerty\r", "\r\n") => 1
# f_p_a_e ("qwertydkjf", "\r\n") => 0
# f_p_a_e ("qwerty\r\n", "\r\n") => <undefined>

# this could maybe be made faster with a computed regex?
# [answer: no; circa Python-2.0, Jan 2001]
# new python:   28961/s
# old python:   18307/s
# re:        12820/s
# regex:     14035/s

def     find_prefix_at_end (haystack, needle):
    l = len(needle) - 1
    while l and not haystack.endswith(needle[:l]):
        l -= 1
    return l

#######################################################################

class ExitNow(Exception): pass

class   Jujucore( object ) :

    def __init__(self) :
        self.map = {}

        if hasattr(select, 'poll') :
            log.info('%s using POLL.', self.__class__.__name__)
            self.poll_func = self.loopByEPoll
        else :
            log.info('%s using SELECT.', self.__class__.__name__)
            self.poll_func = self.loopBySelect
        self.poll_func = self.loopBySelect
        # __init__()

    def addChannel(self, channel) :
        if channel.fileno < 0 :
            # fileno == 0 : means stdin, should be valid
            raise RuntimeError, 'Invalid channel to add (fileno=%d)' % channel.fileno
        self.map[ channel.fileno ] = channel
        log.info('#%d channel check-in (total %d members).', channel.fileno, len(self.map))
        # addChannel()

    def delChannel(self, channel) :
        if channel.fileno < 0 :
            # fileno == 0 : means stdin, should be valid
            raise RuntimeError, 'Invalid channel to delete (fileno=%d)' % channel.fileno

        if channel.fileno in self.map :
            del self.map[ channel.fileno ]
        log.info('#%d channel check-out (total %d members).', channel.fileno, len(self.map))
        # delChannel()

    def loopByEPoll(self, timeout=0.0) :
        map = self.map

        if timeout is not None:
            # timeout is in milliseconds
            timeout = int(timeout*1000)
        pollster = select.poll()
        if map:
            for fd, obj in map.items():
                flags = 0
                if obj.readable():
                    flags |= select.POLLIN | select.POLLPRI
                if obj.writable():
                    flags |= select.POLLOUT
                if flags:
                    # Only check for exceptions if object was either readable
                    # or writable.
                    flags |= select.POLLERR | select.POLLHUP | select.POLLNVAL
                    pollster.register(fd, flags)
            try:
                r = pollster.poll(timeout)
            except select.error, err:
                if err[0] != EINTR:
                    raise
                r = []
            for fd, flags in r:
                obj = map.get(fd)
                if obj is None:
                    continue
                #readwrite(obj, flags)
                try:
                    if flags & (select.POLLIN | select.POLLPRI):
                        obj.handleReadEvent()
                    if flags & select.POLLOUT:
                        obj.handleWriteEvent()
                    if flags & (select.POLLERR | select.POLLHUP | select.POLLNVAL):
                        obj.handleExptionEvent()
                except ExitNow:
                    raise
                except:
                    obj.handleError()
        # loopByEPool()

    def loopBySelect(self, timeout=0.0) :
        map = self.map

        if map:
            r = []; w = []; e = []
            for fd, obj in map.items():
                is_r = obj.readable()
                is_w = obj.writable()
                if is_r:
                    r.append(fd)
                if is_w:
                    w.append(fd)
                if is_r or is_w:
                    e.append(fd)
            if [] == r == w == e:
                time.sleep(timeout)
            else:
                try:
                    r, w, e = select.select(r, w, e, timeout)
                except select.error, err:
                    if err[0] != EINTR:
                        raise
                    else:
                        return

            for fd in r:
                obj = map.get(fd)
                if obj is None:
                    continue
                obj.handleReadEvent()

            for fd in w:
                obj = map.get(fd)
                if obj is None:
                    continue
                obj.handleWriteEvent()

            for fd in e:
                obj = map.get(fd)
                if obj is None:
                    continue
                obj.handleExceptionEvent()
        # loopBySelect()

    def loop(self, timeout=10.0, count=0) :
        _count = 0

        poll_func = self.loopBySelect

        while self.map :
            poll_func(timeout)

            if count > 0 :
                _count += 1
                if _count >= count :
                    break
        # loop()

    # class Jujucore


